#!/usr/bin/python
# by Aaron (ams@bio.aau.dk)
# version 0.5
######################################
# Version History
# v. 0.1 started with the Genbank parser for CDS
# v. 0.2 wrote the arguement parser
# v. 0.3 wrote the proteinsearch function
# v. 0.4 cds parser completed. Added a log file for the cds parser
# v. 0.5 reorganised the functions to use dict for records instead of r/w
# 
# TODO
# organise the output in a dir with the blast result, hits and results and log
# sequence check for protein in proteinsearch
# change main to accept multiple queries in the file or multiple files
#      an compile to one hitlist before CDS parsing
# write the "database" section where the results & queries are saved to file
# maybe then split off into a makedb and a updatedb script
# query files that can be run periodically and return only additions
# implement a summary file of the sequences, some details and date added
#
######################################


import sys 
from Bio import SeqIO
from Bio.Blast import NCBIWWW
from Bio.Blast import NCBIXML
from Bio import Entrez
Entrez.email = 'ams@bio.aau.dk'
import re
import os
from time import localtime, strftime


def parse_cds(gp_records, outdir):
  """ 
  Returns a list of CDS identifiers from a Genbank file  
  
  The input is:
  gp_records  a list containing gp records
  outdir a destination for the parsed cds's
  """
  
  if not os.path.exists(outdir):
      os.mkdir(outdir)
  logfilename = os.path.join(outdir, 'cdsparselog.log')
  savefilename =  os.path.join(outdir, 'cdsparse.txt')
  logfile = open(logfilename, 'w')
  mapfilename =  os.path.join(outdir, 'gp2cds_mapping.txt')
  mapfile = open(mapfilename, 'w')
  
  print 'parsing cds features...'
  cds_not_found = []
  
  for record in SeqIO.parse(gp_records, 'genbank'):
      logfile.write('Parsing: %s\n'%record.id)
      
      # parse out description...
      desc = record.description
      org_start = desc.index('[') - 1
      org_stop = desc.index(']')
      desc = desc[org_start:org_stop]
      pattern = '\[(\w+)\s(\w+)\s(.+)'
      mobj = re.search(pattern, str(desc))
      genus, species, strain = '', '', ''
      if mobj:
          genus = mobj.group(1) 
          species = mobj.group(2) 
          strain = mobj.group(3)
      aa_length = len(record.seq)
      
      # parse out coded_by qualifier...
      coded_by = ''
      for gp_feature in record.features:
          if gp_feature.type == 'CDS':
              coded_by = gp_feature.qualifiers['coded_by']
      if not coded_by:
          logfile.write('Error: CDS not found for %s\n'%record.id)
          cds_not_found.append(record.id)
          continue

      # parse coded_by string: eg. complement(NC_005816.1:4815..5888)
      pattern = '(\w+\.\d+):(\d+)\.\.(\d+)'
      mobj = re.search(pattern, str(coded_by))
      if mobj:
          cds_acc = mobj.group(1) 
          cds_start = mobj.group(2) 
          cds_stop = mobj.group(3) 
          cds_comp = re.search('complement', str(coded_by))
          if cds_comp:
              cds_strand = 'PLUS'
          else:
              cds_strand = 'MINUS'
          logfile.write('CDS parsed...' + cds_acc + ': ' + cds_start
                  + '..' + cds_stop + '\n')
          output = '%s\t\t\t%s\t%s\t%s\t%s\t\t\t%s\t%s\t%s\t%s\n' % (
                  record.id, aa_length, genus, species, strain, cds_acc, 
                  cds_start, cds_stop, cds_strand)
          mapfile.write('{0}'.format(output))
      else:
          logfile.write('Couldn\'t match the CDS \n')
          cds_not_found.append(record.id)
          continue
    
  mapfile.close()
  logfile.close()
  
  return 

# def fetch_cds():
  # logfilename = os.path.join(outdir, 'cdsparselog.log')
  # savefilename =  os.path.join(outdir, 'cdsparse.gb')
  # try:
      # logfile = open(logfilename, 'w')
      # print 'logging to logfile...'
  # except:
      # print 'logfile error'
  # try:
      # save_handle = open(savefilename, 'w')
  # except:
      # print 'savefile error'
  
  # print 'parsing cds features...'
  # cds_not_found = []
  # fetch_count = 0
  # gp2cds_mapping = {}
  
  # # fetch the genbank gene entry
  # logfile.write('sent Entrez request...  /n')
  # try:
      # handle = Entrez.efetch(db='nuccore', id=x_id, seq_start=x_seq_start,
              # seq_stop=x_seq_stop, strand=x_strand, retmode='gb')
  # except:
      # logfile.write('Entrez fetch error!\n')
      # cds_not_found.append(record.id)
      # continue
  # logfile.write('Entrez fetch... [ok]\n')
  # try:
      # save_handle.write(handle.read())
      # logfile.write('CDS appended to %s... [ok]\n'%savefilename)
      # gp2cds_mapping[record.id] = x_id
  # except:
      # logfile.write('Save error for entry: %s \n'%record.id)
      # continue
  # fetch_count += 1
  
  # report = 'CDS fetch completed... [ok]\n'
  # report = report + '{0} records fetched\n'.format(fetch_count)
  # report = report + '{0} cds not retrieved\n'.format(len(cds_not_found))
  # report = report + 'CDS not found for: ' + '\n'.join(cds_not_found)
  # print report
  # logfile.write(report)
  # logfile.close()

def main(): 
  
  infilename = sys.argv[1]
  outdir = 'cdsfetch'
  if not os.path.exists(outdir):
      os.mkdir(outdir)
  
  gp_records = open(infilename)

  cds_records = parse_cds(gp_records, outdir)
  

  
if __name__ == '__main__':
  main()


